import numpy as np
from pymilvus import connections, Collection, FieldSchema, CollectionSchema, DataType, utility

# 连接Milvus
connections.connect("default", host="localhost", port="19530")

# 定义Collection
fields = [
    FieldSchema(name="id", dtype=DataType.INT64, is_primary=True),
    FieldSchema(name="embedding", dtype=DataType.FLOAT_VECTOR, dim=384),  # BERT向量维度
    FieldSchema(name="text", dtype=DataType.VARCHAR, max_length=65535),
    FieldSchema(name="category", dtype=DataType.VARCHAR, max_length=255)
]
collection_name = "documents"
schema = CollectionSchema(fields, "文档语义搜索系统")


# 检查集合是否存在
if utility.has_collection(collection_name):
    # 如果存在，直接加载集合
    collection = Collection(name=collection_name)
    print(f"Collection {collection_name} already exists. Using existing collection.")

else:
    # 如果不存在，创建新集合
    collection = Collection(name=collection_name, schema=schema)
    print(f"Collection {collection_name} created.")
    # 创建索引
    index_params = {
        "metric_type": "COSINE",  # 余弦相似度
        "index_type": "HNSW",  # 使用HNSW索引
        "params": {"M": 16, "efConstruction": 200}
    }
    collection.create_index("embedding", index_params)




